5.5 读写文件

1. 加载和保存张量

对于单个张量，我们可以直接调用 load 和 save 函数分别读写它们：

import torch
from torch import nn
from torch.nn import functional as F

x = torch.arange(4)
torch.save(x, 'x-file')

我们现在可以将存储在文件中的数据读回内存

x2 = torch.load('x-file')
x2

tensor([0, 1, 2, 3])

对于多个张量，我们可以将张量列表或者张量字典存到文件，也可以反过来读取存到为张量列表或者张量字典：

y = torch.zeros(4)
torch.save([x, y],'x-files')
x2, y2 = torch.load('x-files')
(x2, y2) # (tensor([0, 1, 2, 3]), tensor([0., 0., 0., 0.]))


mydict = {'x': x, 'y': y}
torch.save(mydict, 'mydict')
mydict2 = torch.load('mydict')
mydict2 # {'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}

2. 将模型保存到文件

直接将 module 类保存到文件是一件困难的事情，如果要确定一个网络，重要的是两个要素：

• 模型架构（模型定义的代码）
• 模型参数大小（一堆张量）

所以，模型架构是一段代码，这段代码自己保留，模型参数的参数可以保留到文件中，以 MLP 为例：

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)

    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))

net = MLP()

现在将参数字典保留到文件：

torch.save(net.state_dict(), 'mlp.params')

为了恢复模型，我们先读取模型架构，再读取模型参数：

clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))
clone.eval()

MLP(
  (hidden): Linear(in_features=20, out_features=256, bias=True)
  (output): Linear(in_features=256, out_features=10, bias=True)
)

下面我们验证是不是同一个模型：

X = torch.randn(size=(2, 20))

Y = net(X)
Y_clone = clone(X)

Y_clone == Y

tensor([[True, True, True, True, True, True, True, True, True, True],
        [True, True, True, True, True, True, True, True, True, True]])

参考文献